import numpy as np
from collections import namedtuple
from env import Env

class ValueTable:
    def __init__(self, env):
        self._values = np.zeros(env.state_space)

    def update(self, s, value):
        self._values[s] = value

    def get(self, state):
        if type(state) == list:
            # loop get
            res = [self._values[s] for s in state]
            return res
        elif type(state) == tuple:
            # return directly
            return self._values[state]




Pi = namedtuple('Pi', 'act, prob')


class Policies:
    def __init__(self, env: Env):
        self._actions = env.actions
        self._default_policy = [1 / env.action_space] * env.action_space
        self._policies = dict.fromkeys(env.states, Pi(self._actions, self._default_policy))

    def sample(self, state):
        if self._policies.get(state, None) is None:
            self._policies[state] = Pi(self._actions, self._default_policy)

        policy = self._policies[state]
        return np.random.choice(policy.act, p=policy.prob)

    def retrieve(self, state):
        return self._policies[state].prob

    def update(self, state, policy):
        self._policies[state] = self._policies[state]._replace(prob=policy)